package com.mezz.bones.framework.core.thread;

import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.Callable;
import java.util.concurrent.Future;

/**
 * 自定义线程池 增强 异常输出信息
 */
@Slf4j
public class WrappedThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

    @Override
    public void execute(Runnable task) {
        super.execute(createWrappedRunnable(task));
    }

    @Override
    public void execute(Runnable task, long startTimeout) {
        super.execute(createWrappedRunnable(task), startTimeout);
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(createWrappedCallable(task));
    }

    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(createWrappedRunnable(task));
    }

    private Runnable createWrappedRunnable(Runnable task) {
        return () -> {
            try {
                task.run();
            } catch (Exception e) {
                handleUncaughtException(e, task);
                throw e;
            }
        };
    }

    private <T> Callable<T> createWrappedCallable(Callable<T> task) {
        return () -> {
            try {
                return task.call();
            } catch (Exception e) {
                handleUncaughtException(e, task);
                throw e; // Rethrow to ensure proper exception handling
            }
        };
    }

    private void handleUncaughtException(Exception e, Object task) {
        log.error("Uncaught exception in task: {}", task, e);
    }
}
